from langchain_core.messages import AIMessage, HumanMessage
from  langgraph.graph import StateGraph
from  langgraph.graph import MessageGraph,MessagesState
from typing import Any

# 创建一个工作流
# MessagesState相当于构建了一个状态实体
workflow = StateGraph(MessagesState)

# 定义大模型处理的消息状态
def chat_bot(state:MessagesState):
    return {"messages": [AIMessage(content="你好，我是OpenAI开发的聊天机器人")]}

# 定义并行执行的函数
def parallel1(state: MessagesState, config: dict) -> Any:
    print("并行1: ", state)
    return {"messages": [HumanMessage(content="这是并行1函数")]}


def parallel2(state: MessagesState, config: dict) -> Any:
    print("并行2: ", state)
    return {"messages": [HumanMessage(content="这是并行2函数")]}


def chat_end(state: MessagesState, config: dict) -> Any:
    print("聊天结束: ", state)
    return {"messages": [HumanMessage(content="这是聊天结束函数")]}
# 添加节点
workflow.add_node('chat_bot',chat_bot)
workflow.add_node('chat_end',chat_end)
workflow.add_node('parallel1',chat_end)
workflow.add_node('parallel2',chat_end)

workflow.set_entry_point('chat_bot')
workflow.set_finish_point('chat_end')
workflow.add_edge("chat_bot", "parallel1")
workflow.add_edge("chat_bot", "parallel2")
workflow.add_edge("parallel2", "chat_end")

graph = workflow.compile()
print(graph.invoke({"messages": [HumanMessage(content="你好，你是")]}))